
#include "glodef.h"
#include "errlog.h"
#include "vecs.h"
#include "grid3d.h"
#include "fe.h"
#include "fe_matr.h"
// these are the standard C library files
#include <stdlib.h>
#include <memory.h>
// these are the PETSc files
#include <petsc.h>
#include <petscmat.h>
#include <petscvec.h>
#include <petscksp.h>


#include "vecs_petsc.h"


#define CHKERR_B(n) if (n) {n = PetscError(__LINE__, __FUNCT__, __FILE__, __SDIR__, n, 0, " ");error_msg("%s: petsc error detected", __func__);}

#define USE_CONDENSATION 1

#if !defined(USE_CONDENSATION)
#define USE_CONDENSATION    0
#endif

/*
 *  ATTENTION: Condensation only works for p1-bubble element.
 *  condense_data is an array of FLOATs with size 9*grid->NEL.
 *  Denote lM[i][j] the local matrix of element el.  Then condense_data
 *  contains the following data:
 *      condense_data[el*9 + j] == lM[4][j], for j = 0, 1, 2, 3, 4
 *      condense_data[el*9 + 5 + i] == lM[i][4] / lM[4][4], for i = 0, 1, 2, 3
 *  The first set of numbers is needed to compute the values of the bubble
 *  degrees of freedom in post-processing.  The second set are used to modify
 *  the right hand side of the system.
 */


typedef struct _fe_matr_t
{
    unsigned el_type;
    grid3d_t *grid;
    local_matr_t *local_matr;
    get_eo_t *get_eo;
    Mat matr;
    Mat matr_bc;
    KSP ksp;
    void *condense_data;

    struct fe_matr_els
    {
        int local_NEL;
        int *local_elems;
    } *elem_data;

    struct fe_matr_ss
    {
        int size;
        int *sm_d_nnz;
        int *sm_o_nnz;
        int **sm_nzcols;
        int *bm_d_nnz;
        int *bm_o_nnz;
        int **bm_nzcols;
    } *structure;
} fe_matr_t;

unsigned get_el_type(fe_matr_p M)
{
    return M->el_type;
}

static int check_dimensions(Mat sm, Mat bm)
{
    int dim, dim1, dim2, dim3, ierr;
    ierr=MatGetSize(sm, &dim, &dim1);
    CHKERR_B(ierr);
    ierr=MatGetSize(bm, &dim2, &dim3);
    CHKERR_B(ierr);
    assert((dim==dim1)&&(dim1==dim2)&&(dim2==dim3));
    return dim;
}

static void fe_matr_ss_alloc(fe_matr_p matr, int rows)
{
    struct fe_matr_ss *S=calloc(1, sizeof(*S));
    if(S==NULL) goto alloc_error;
    int low, high;
    MatGetOwnershipRange(matr->matr, &low, &high);
    S->size=high-low;

    S->sm_d_nnz=calloc(S->size, sizeof(int));
    S->sm_o_nnz=calloc(S->size, sizeof(int));
    S->sm_nzcols=calloc(S->size, sizeof(int *));
    S->bm_d_nnz=calloc(S->size, sizeof(int));
    S->bm_o_nnz=calloc(S->size, sizeof(int));
    S->bm_nzcols=calloc(S->size, sizeof(int *));
    if(S->sm_d_nnz==NULL||S->sm_o_nnz==NULL||S->sm_nzcols==NULL
        ||S->bm_d_nnz==NULL||S->bm_o_nnz==NULL||S->bm_nzcols==NULL)
        goto alloc_error;
    matr->structure=S;
    return;
alloc_error:
    error_msg("%s: memory allocation error.\n", __func__);
}

static void fe_matr_ss_free(fe_matr_p matr)
{
    if(matr->structure)
    {
        struct fe_matr_ss *S=matr->structure;
        for(int i=0; i<S->size; i++)
            free(S->sm_nzcols[i]);
        free(S->sm_nzcols);
        free(S->sm_d_nnz);
        free(S->sm_o_nnz);
        for(int i=0; i<S->size; i++)
            free(S->bm_nzcols[i]);
        free(S->bm_nzcols);
        free(S->bm_d_nnz);
        free(S->bm_o_nnz);
        free(matr->structure);
    }
}

enum
{
    sm_add, sm_ins, bm_add, bm_ins
};

static void fe_matr_set_value(fe_matr_p mat, int row, int col, double val, int mode)
{
    int ierr;
    int low, high;
    MatGetOwnershipRange(mat->matr, &low, &high);
    if((row<low)||(row>=high)) return;
    if(!mat->structure)
    {
        switch(mode)
        {
            case sm_add:
                ierr=MatSetValues(mat->matr, 1, &row, 1, &col, &val, ADD_VALUES);
                CHKERR_B(ierr);
                break;
            case sm_ins:
                ierr=MatSetValues(mat->matr, 1, &row, 1, &col, &val, INSERT_VALUES);
                CHKERR_B(ierr);
                break;
            case bm_add:
                ierr=MatSetValues(mat->matr_bc, 1, &row, 1, &col, &val, ADD_VALUES);
                CHKERR_B(ierr);
                break;
            case bm_ins:
                ierr=MatSetValues(mat->matr_bc, 1, &row, 1, &col, &val, INSERT_VALUES);
                CHKERR_B(ierr);
                break;
            default:
                error_msg("%s: invalid `mode' parameter.\n", __func__);
        }
        return;
    }
    else /* just building structure, not really assembly at this time */
    {
        struct fe_matr_ss *S=mat->structure;
        int *d_nnz, *o_nnz, **nzcols;
        row=row-low;
        switch(mode)
        {
            case sm_add:
            case sm_ins:
                d_nnz=S->sm_d_nnz;
                o_nnz=S->sm_o_nnz;
                nzcols=S->sm_nzcols;
                break;
            case bm_add:
            case bm_ins:
                d_nnz=S->bm_d_nnz;
                o_nnz=S->bm_o_nnz;
                nzcols=S->bm_nzcols;
                break;
            default:
                error_msg("%s: invalid `mode' parameter.\n", __func__);
        }
        int bar=d_nnz[row]+o_nnz[row];
        for(int i=0; i<bar; i++)
            if(nzcols[row][i]==col)
                return;
        if((col<low)||(col>=high))
            o_nnz[row]++;
        else
            d_nnz[row]++;
        bar++;
        int *foo=realloc(nzcols[row], bar*sizeof(int));
        if(foo==NULL)
            error_msg("%s: memory re-allocation error.\n", __func__);
        nzcols[row]=foo;
        nzcols[row][bar-1]=col;
    }
}

static void entry_oper(fe_matr_p m, int ni, int nj,
    unsigned eoi, unsigned eoj, double val, int add_insert)
{
    const double one=1.0;
    if((eoi==eo_ignore)||(eoj==eo_ignore))
        return;
    if((eoi==eo_bc)||(eoj==eo_bc))
    {
        if(ni==nj)
            // SM_SetEntry(sm, ni, nj, 1.0);
            if(add_insert==INSERT_VALUES)
                fe_matr_set_value(m, ni, nj, one, sm_ins);
        // MatSetValues(sm, 1, &ni, 1, &nj, &one, INSERT_VALUES);
        switch(eoi)
        {
            case eo_add:
                // SM_AddToEntry(bm, ni, nj, val);
                if(add_insert==ADD_VALUES)
                    fe_matr_set_value(m, ni, nj, val, bm_add);
                // MatSetValues(bm, 1, &ni, 1, &nj, &val, ADD_VALUES);
                break;
            case eo_set:
                // SM_SetEntry(bm, ni, nj, val);
                if(add_insert==INSERT_VALUES)
                    fe_matr_set_value(m, ni, nj, val, bm_ins);
                // MatSetValues(bm, 1, &ni, 1, &nj, &val, INSERT_VALUES);
                break;
            default: break;
        }
        return;
    }
    switch(eoi)
    {
        case eo_add:
            switch(eoj)
            {
                case eo_add:
                    // if((!sm->Symmetric)||(ni<=nj))
                    //     SM_AddToEntry(sm, ni, nj, val);
                    if(add_insert==ADD_VALUES)
                        fe_matr_set_value(m, ni, nj, val, sm_add);
                    // MatSetValues(sm, 1, &ni, 1, &nj, &val, ADD_VALUES);
                    return;
                case eo_set:
                    goto err_setadd;
            }
        case eo_set:
            switch(eoj)
            {
                case eo_add:
                    goto err_setadd;
                case eo_set:
                    // if((!sm->Symmetric)||(ni<=nj))
                    //     SM_SetEntry(sm, ni, nj, val);
                    if(add_insert==INSERT_VALUES)
                        fe_matr_set_value(m, ni, nj, val, sm_ins);
                    // MatSetValues(sm, 1, &ni, 1, &nj, &val, INSERT_VALUES);
                    return;
            }
    }
err_setadd:
    error_msg("fe_matr_assemble: ambiguous set/add entry operation.\n");
}

static void fe_matr_ss_create(fe_matr_p matr)
{
    grid3d_t *grid=matr->grid;
    int add_insert=ADD_VALUES;
    while(1)
    {
        for(int el=0; el<grid->NEL; el++)
        {
            unsigned *eo;
            int *ln, nn, i, j;
            nn=fe_dofs(grid, el, matr->el_type, &ln);
            matr->get_eo(matr->grid, matr->el_type, ln, &eo);
            // do condensation
            if(matr->el_type==fe_bub&&matr->condense_data!=NULL&&USE_CONDENSATION)
            {
                nn--;
            }
            else if(matr->el_type==fe_bub_vec&&matr->condense_data!=NULL&&USE_CONDENSATION)
            {
                int k, j;
                int NN=grid->NN;
                int NEL=grid->NEL;
                static int foo[15];

                for(k=0; k<3; k++)
                {
                    for(j=0; j<4; j++)
                        foo[5*k+j]=NN*k+ln[j];
                    foo[5*k+4]=3*NN+k*NEL+el;
                }
                ln=foo;
                eo[4]=eo[9]=eo[14]=eo_ignore;
            }

            for(i=0; i<nn; i++)
                for(j=0; j<nn; j++)
                    entry_oper(matr, ln[i], ln[j], eo[i], eo[j], 0.0, add_insert);

        }
        if(add_insert==INSERT_VALUES)
            break;
        add_insert=INSERT_VALUES;
    }
}

fe_matr_p fe_matr_alloc(unsigned type, local_matr_t *lm, get_eo_t *get_eo,
    grid3d_t *grid, int symmetric, char *mattyp)
{
    fe_matr_p newmat;
    int ndof, ierr;
    LOG_MPI_WTIME_START;

    newmat=calloc(1, sizeof(fe_matr_t));
    if(newmat==NULL)
        goto alloc_err;
    newmat->el_type=type;
    newmat->grid=grid;
    newmat->local_matr=lm;
    newmat->get_eo=get_eo;

    if(type==fe_bub&&USE_CONDENSATION)
    {
        double *cd=malloc(9*grid->NEL*sizeof(double));
        if(cd==NULL)
            goto alloc_err;
        newmat->condense_data=cd;
        ndof=grid->NN;
    }
    else if(type==fe_bub_vec&&USE_CONDENSATION)
    {
        double *cd=malloc(90*grid->NEL*sizeof(double));
        if(cd==NULL)
            goto alloc_err;
        newmat->condense_data=cd;
        ndof=3*grid->NN;
    }
    else
    {
        newmat->condense_data=NULL;
        ndof=fe_ndof(grid, type);
    }
    ierr=MatCreate(PETSC_COMM_WORLD, &newmat->matr);
    CHKERR_B(ierr);
    ierr=MatSetSizes(newmat->matr, PETSC_DECIDE, PETSC_DECIDE, ndof, ndof);
    CHKERR_B(ierr);
    ierr=MatSetType(newmat->matr, MATAIJ);
    CHKERR_B(ierr);

    ierr=MatCreate(PETSC_COMM_WORLD, &newmat->matr_bc);
    CHKERR_B(ierr);
    ierr=MatSetSizes(newmat->matr_bc, PETSC_DECIDE, PETSC_DECIDE, ndof, ndof);
    CHKERR_B(ierr);
    ierr=MatSetType(newmat->matr_bc, MATAIJ);
    CHKERR_B(ierr);

    fe_matr_ss_alloc(newmat, ndof);
    fe_matr_ss_create(newmat);

    if(mpi_num_procs>1)
    {
        ierr=MatMPIAIJSetPreallocation(newmat->matr, PETSC_DEFAULT, newmat->structure->sm_d_nnz, PETSC_DEFAULT, newmat->structure->sm_o_nnz);
        CHKERR_B(ierr);
        ierr=MatMPIAIJSetPreallocation(newmat->matr_bc, PETSC_DEFAULT, newmat->structure->bm_d_nnz, PETSC_DEFAULT, newmat->structure->bm_o_nnz);
        CHKERR_B(ierr);
    }
    else
    {
        ierr=MatSeqAIJSetPreallocation(newmat->matr, PETSC_DEFAULT, newmat->structure->sm_d_nnz);
        CHKERR_B(ierr);
        ierr=MatSeqAIJSetPreallocation(newmat->matr_bc, PETSC_DEFAULT, newmat->structure->bm_d_nnz);
        CHKERR_B(ierr);
    }

    fe_matr_ss_free(newmat);
    newmat->structure=NULL;

    /* make a list of elements that contain the dofs that we need to update */
    /* {
        char myfname[1000];
        sprintf(myfname, "elems.%d.txt",mpi_my_id);
        FILE *F = fopen(myfname, "wt");

        int lo, hi, el;
        int nn, *ln, i;
        MatGetOwnershipRange(newmat->matr, &lo, &hi);
        if(F) fprintf(F, "lo=%d, hi=%d\n", lo, hi);
        struct fe_matr_els *edata;
        edata=calloc(1, sizeof(struct fe_matr_els));
        if(edata==NULL)
            error_msg("%s: cannot allocate local elem_data\n", __func__);
        else
            newmat->elem_data=edata;
        for(el=0; el<grid->NEL; el++)
        {
            nn=fe_dofs(grid, el, newmat->el_type, &ln);
            for(i=0; i<nn; i++)
                if((ln[i]>=lo)&&(ln[i]<hi))
                    break;
            if(i<nn)
            {
                int *bar=NULL, foo=edata->local_NEL;
                foo++;
                bar=realloc(edata->local_elems, foo*sizeof(int));
                if(bar==NULL)
                    error_msg("%s: cannot allocate local elem_data II \n", __func__);
                bar[foo-1]=el;
                edata->local_NEL=foo;
                edata->local_elems=bar;
                if(F)
                {
                    fprintf(F, "elem %d:", el);
                    for(i=0;i<nn;i++)
                        fprintf(F, "  %d", ln[i]);
                    fprintf(F, "\n");
                }
            }
        }
        if(F) fclose(F);
    }*/

    KSP ksp;
    ierr=KSPCreate(PETSC_COMM_WORLD, &ksp);
    CHKERR_B(ierr);
    ierr=KSPSetOperators(ksp, newmat->matr, newmat->matr, DIFFERENT_NONZERO_PATTERN);
    CHKERR_B(ierr);
    ierr=KSPSetFromOptions(ksp);
    CHKERR_B(ierr);
    ierr=KSPSetInitialGuessNonzero(ksp, PETSC_TRUE);
    CHKERR_B(ierr);
    ierr=KSPSetTolerances(ksp, LS_ACCURACY, LS_ACCURACY, PETSC_DEFAULT, PETSC_DEFAULT);
    CHKERR_B(ierr);
    if(symmetric)
    {
        MatSetOption(newmat->matr, MAT_SYMMETRIC, PETSC_TRUE);
        MatSetOption(newmat->matr, MAT_SYMMETRY_ETERNAL, PETSC_TRUE);
        ierr=KSPSetType(ksp, KSPCG);
        CHKERR_B(ierr);
    }
    else
    {
        ierr=KSPSetType(ksp, KSPBCGS);
        CHKERR_B(ierr);
    }
    newmat->ksp=ksp;

    LOG_MPI_WTIME_STOP;

    return newmat;

alloc_err:
    error_msg("%s: memory allocation error.\n", __func__);
}

void fe_matr_free(fe_matr_p matr)
{
    if(matr)
    {
        int ierr;
        ierr=MatDestroy(matr->matr);
        CHKERR_B(ierr);
        ierr=MatDestroy(matr->matr_bc);
        CHKERR_B(ierr);
        ierr=KSPDestroy(matr->ksp);
        CHKERR_B(ierr);
        if(matr->elem_data)
        {
            if(matr->elem_data->local_elems)
                free(matr->elem_data->local_elems);
            free(matr->elem_data);
        }
        if(matr->condense_data)
            free(matr->condense_data);
        fe_matr_ss_free(matr);
        free(matr);
    }
}

void fe_matr_assemble(const grid3d_t *grid, int nmatr, fe_matr_p* matr)
{
    int el, mi, ierr;
    LOG_MPI_WTIME_START;
    int add_insert;

    PrintLog("Assembling %u matrices of dimensions ", nmatr);

    for(mi=0; mi<nmatr; mi++)
    {
        int m, n;
        ierr=MatZeroEntries(matr[mi]->matr);
        CHKERR_B(ierr);
        ierr=MatZeroEntries(matr[mi]->matr_bc);
        CHKERR_B(ierr);
        ierr=MatGetSize(matr[mi]->matr, &m, &n);
        CHKERR_B(ierr);
        PrintLog("%u ", m);
    }

    add_insert=ADD_VALUES;
    while(1)
    {
        for(mi=0; mi<nmatr; ++mi)
            /* if(matr[mi]) */
        {
            fe_matr_p m=matr[mi];
            int lNEL, *lelems;
            /* if(m->elem_data)
            {
                lNEL=m->elem_data->local_NEL;
                lelems=m->elem_data->local_elems;
            }
            else */
            {
                lNEL=grid->NEL;
                lelems=NULL;
            }
            for(int ii=0; ii<lNEL; ii++)
            {
                el=ii; // lelems?lelems[ii]:ii;
                double **lM;
                unsigned *eo;
                int *ln, nn, i, j;
                m->local_matr(grid, el, m->el_type, &lM);
                nn=fe_dofs(grid, el, m->el_type, &ln);
                m->get_eo(grid, m->el_type, ln, &eo);
                /* do condensation */
                //                if(add_insert == ADD_VALUES)
                {
                    if(m->el_type==fe_bub&&m->condense_data!=NULL
                        &&USE_CONDENSATION)
                    {
                        double *CD=m->condense_data;
                        for(i=0; i<4; i++)
                        {
                            double CC=lM[i][4]/lM[4][4];
                            if(eo[i]==eo_bc)
                                CC=0.0;
                            else
                                for(j=0; j<4; j++)
                                    lM[i][j]-=CC*lM[4][j];
                            CD[9*el+i]=lM[4][i];
                            CD[9*el+5+i]=CC;
                        }
                        CD[9*el+4]=lM[4][4];
                        nn--;
                    }
                    else if(m->el_type==fe_bub_vec&&m->condense_data!=NULL
                        &&USE_CONDENSATION)
                    {
                        int k;
                        double *CD=((double *)m->condense_data)+90*el;
                        /* save coefficients of Gauss elimination of cols 4, 9, 14 */
                        for(k=0; k<3; k++)
                            for(i=0; i<15; i++)
                            {
                                double Coef;
                                if(i==4+5*k) /* no elimination in condensation rows */
                                    Coef=0.0;
                                else if(eo[i]==eo_bc) /* no elimination in rows of Dirchlet BC */
                                    Coef=0.0;
                                else /* Gaussian elimination */
                                {
                                    Coef=lM[i][4+5*k]/lM[4+5*k][4+5*k];
                                    for(j=0; j<15; j++)
                                        lM[i][j]-=Coef*lM[4+5*k][j];
                                }
                                CD[(3+k)*15+i]=Coef;
                            }
                        /* save rows 4, 9, and 14 */
                        for(k=0; k<3; k++)
                            for(i=0; i<15; i++)
                                CD[15*k+i]=lM[5*k+4][i];

                        int NN=grid->NN;
                        int NEL=grid->NEL;
                        int foo[15];

                        for(k=0; k<3; k++)
                        {
                            for(j=0; j<4; j++)
                                foo[5*k+j]=NN*k+ln[j];
                            foo[5*k+4]=3*NN+k*NEL+el;
                        }
                        ln=foo;
                        eo[4]=eo[9]=eo[14]=eo_ignore;
                    }
                }

                for(i=0; i<nn; i++) for(j=0; j<nn; j++)
                    {
                        entry_oper(m, ln[i], ln[j], eo[i], eo[j], lM[i][j],
                            add_insert);
                    }
            }
        }

        // break;
        if(add_insert==INSERT_VALUES)
            break;
        add_insert=INSERT_VALUES;
        for(mi=0; mi<nmatr; mi++)
        {
            ierr=MatAssemblyBegin(matr[mi]->matr, MAT_FLUSH_ASSEMBLY);
            CHKERR_B(ierr);
            ierr=MatAssemblyEnd(matr[mi]->matr, MAT_FLUSH_ASSEMBLY);
            CHKERR_B(ierr);
            ierr=MatAssemblyBegin(matr[mi]->matr_bc, MAT_FLUSH_ASSEMBLY);
            CHKERR_B(ierr);
            ierr=MatAssemblyEnd(matr[mi]->matr_bc, MAT_FLUSH_ASSEMBLY);
            CHKERR_B(ierr);
        }
    }

    LOG_MPI_WTIME_STOP;

    double tm=MPI_Wtime();
    PrintLog("PETSC Assemble: ");
    for(mi=0; mi<nmatr; mi++)
    {
        ierr=MatAssemblyBegin(matr[mi]->matr, MAT_FINAL_ASSEMBLY);
        CHKERR_B(ierr);
        ierr=MatAssemblyEnd(matr[mi]->matr, MAT_FINAL_ASSEMBLY);
        CHKERR_B(ierr);
        ierr=MatAssemblyBegin(matr[mi]->matr_bc, MAT_FINAL_ASSEMBLY);
        CHKERR_B(ierr);
        ierr=MatAssemblyEnd(matr[mi]->matr_bc, MAT_FINAL_ASSEMBLY);
        CHKERR_B(ierr);
    }
    PrintLog("    CPU time  %-6.3fsec.\n", (MPI_Wtime()-tm));
}


/***********************************************************************************************/
/*  Solution of linear system by converting VECSEQ vectors on SELF to VECMPI vectors on WORLD  */

/***********************************************************************************************/

static int mat_solve_bm(KSP ksp, Mat sm, Mat bm, Vec F, Vec X, int dim)
{
    int ret=TRUE;
    int iters, ierr;
    double res;
    KSPConvergedReason reason;
    Vec scalX, scalF, scalT;

    ierr=VecCreate(PETSC_COMM_WORLD, &scalX);
    CHKERR_B(ierr);
    ierr=VecSetSizes(scalX, PETSC_DECIDE, dim);
    CHKERR_B(ierr);
    ierr=VecSetType(scalX, VECMPI);
    CHKERR_B(ierr);

    ierr=VecDuplicate(scalX, &scalF);
    CHKERR_B(ierr);
    ierr=VecDuplicate(scalX, &scalT);
    CHKERR_B(ierr);

    int i, low, high, size;
    VecGetOwnershipRange(scalX, &low, &high);
    size=high-low;

    for(i=low; i<high; i++)
    {
        double val;

        VecGetValues(X, 1, &i, &val);
        VecSetValues(scalX, 1, &i, &val, INSERT_VALUES);

        VecGetValues(F, 1, &i, &val);
        VecSetValues(scalF, 1, &i, &val, INSERT_VALUES);
    }

    VecAssemblyBegin(scalX);
    VecAssemblyEnd(scalX);
    VecAssemblyBegin(scalF);
    VecAssemblyEnd(scalF);

    ierr=MatMult(bm, scalF, scalT);
    CHKERR_B(ierr);
    ierr=VecAXPY(scalF, -1.0, scalT);
    CHKERR_B(ierr);

    double tm=MPI_Wtime();
    // PrintLog("%s is here\n", __func__ );
    PrintLog(" -> ");
    // ierr=KSPSetOperators(ksp, sm, sm, DIFFERENT_NONZERO_PATTERN);
    // CHKERR_B(ierr);
    ierr=KSPSolve(ksp, scalF, scalX);
    CHKERR_B(ierr);
    ierr=KSPGetIterationNumber(ksp, &iters);
    CHKERR_B(ierr);
    ierr=KSPGetResidualNorm(ksp, &res);
    CHKERR_B(ierr);
    ierr=KSPGetConvergedReason(ksp, &reason);
    CHKERR_B(ierr);
    PrintLog(" ksp:  dim=%-7u  iters=%-6u  res=%-8.2e  reason=%-3d time=%-.3fsec.\n",
        dim, iters, res, reason, (MPI_Wtime()-tm));
    if(reason<0)
    {
        warning_msg("KSP did not converge with reason code %d\n", reason);
        ret=FALSE;
    }

    double *scalx_array, *x_array;
    ierr=VecGetArray(scalX, &scalx_array);
    CHKERR_B(ierr);
    ierr=VecGetArray(X, &x_array);
    CHKERR_B(ierr);
    int tag;
    ierr=PetscObjectGetNewTag((PetscObject)scalX, &tag);
    CHKERR_B(ierr);
    PetscMap *map=scalX->map;
    memcpy(x_array+low, scalx_array, size*sizeof(double));
    for(i=0; i<mpi_num_procs; i++)
    {
        int first, last, total;
        first=map->range[i];
        last=map->range[i+1];
        total=last-first;

        MPI_Bcast(x_array+first, total, MPI_DOUBLE, i, PETSC_COMM_WORLD);
    }

    ierr=VecRestoreArray(scalX, &scalx_array);
    CHKERR_B(ierr);
    ierr=VecRestoreArray(X, &x_array);
    CHKERR_B(ierr);

    ierr=VecAssemblyBegin(X);
    CHKERR_B(ierr);
    ierr=VecAssemblyEnd(X);
    CHKERR_B(ierr);

    ierr=VecDestroy(scalX);
    CHKERR_B(ierr);
    ierr=VecDestroy(scalF);
    CHKERR_B(ierr);
    ierr=VecDestroy(scalT);
    CHKERR_B(ierr);

    return ret;
}

int fe_matr_solve_scalar(const fe_matr_p matr, scalar X, scalar F, scalar T, double TOL)
{
    int i;
    if(matr==NULL)
        error_msg("%s: NULL matrix passed.\n", __func__);

    KSP ksp=matr->ksp;
    Mat sm=matr->matr;
    Mat bm=matr->matr_bc;
    int dim=check_dimensions(sm, bm);

    /* assert((X->N==dim) && (F->N==dim) && T->N==dim); */

    LOG_MPI_WTIME_START;
    PrintLog("Solving scalar equation of dim = %u\n", dim);

    /* pre-processing for condensation */
    if(matr->el_type==fe_bub_vec&&matr->condense_data!=NULL&&USE_CONDENSATION)
    {
        int el, *ln, k;
        double *CD=(double *)matr->condense_data;

        use_scalar(F)
        use_scalar(T)
        use_scalar(X)

        for(el=0; el<matr->grid->NEL; el++, CD+=90)
        {
            fe_dofs(matr->grid, el, fe_bub_vec, &ln);
            for(k=0; k<3; k++)
                for(i=0; i<15; i++)
                    FF[ln[i]]-=CD[(3+k)*15+i]*FF[ln[4+5*k]];
        }

        int NN=matr->grid->NN;
        int NEL=matr->grid->NEL;
        int NBUB=NN+NEL;
        memcpy(TF, FF, F->Size*sizeof(double));
        for(k=0; k<3; k++)
        {
            for(i=0; i<NN; i++) FF[k*NN+i]=TF[k*NBUB+i];
            for(i=0; i<NEL; i++) FF[3*NN+k*NEL+i]=TF[k*NBUB+NN+i];
        }

        done_scalar(X)
        done_scalar(T)
        done_scalar(F)
    }

    clean_scal(T);


    int ret=mat_solve_bm(ksp, sm, bm, F->VecF, X->VecF, dim);

    /* post-processing for condensation */
    if(matr->el_type==fe_bub_vec&&matr->condense_data!=NULL&&USE_CONDENSATION)
    {
        int k, NN=matr->grid->NN;
        int NEL=matr->grid->NEL;
        int NBUB=NN+NEL;
        copy_scal(X, T);
        use_scalar(F)
        use_scalar(T)
        use_scalar(X)
        for(k=0; k<3; k++)
        {
            for(i=0; i<NN; i++) XF[k*NBUB+i]=TF[k*NN+i];
            for(i=0; i<NEL; i++) XF[k*NBUB+NN+i]=TF[3*NN+k*NEL+i];
        }

        int el, *ln;
        double *CD=(double *)matr->condense_data;
        for(el=0; el<matr->grid->NEL; el++, CD+=90)
        {
            fe_dofs(matr->grid, el, fe_bub_vec, &ln);
            for(k=0; k<3; k++)
            {
                XF[ln[5*k+4]]=FF[3*NN+k*NEL+el];
                for(i=0; i<15; i++)
                    XF[ln[5*k+4]]-=((i==5*k+4)?0.0:CD[k*15+i])*XF[ln[i]];
                XF[ln[5*k+4]]/=CD[k*15+5*k+4];
            }
        }
        done_scalar(X)
        done_scalar(T)
        done_scalar(F)
    }


    LOG_MPI_WTIME_STOP;
    return ret;
}

int fe_matr_solve_vector(const fe_matr_p matr, vector X, vector F, vector T, double TOL)
{
    int ret;
    int i;
    if(matr==NULL)
        error_msg("%s: NULL matrix passed.\n", __func__);

    KSP ksp=matr->ksp;
    Mat sm=matr->matr;
    Mat bm=matr->matr_bc;
    int dim=check_dimensions(sm, bm);

    LOG_MPI_WTIME_START;
    PrintLog("Solving vector equation of dim = %u\n", dim);


    /* pre-processing for condensation */
    if(matr->el_type==fe_bub&&matr->condense_data!=NULL&&USE_CONDENSATION)
    {
        int el, *ln;
        double *CD=matr->condense_data;
        use_vector(F)
            // use_vector(T)
            // use_vector(X)
        for(el=0; el<matr->grid->NEL; el++)
        {
            fe_dofs(matr->grid, el, fe_bub, &ln);
            for(i=0; i<4; i++)
            {
                FX[ln[i]]-=CD[9*el+5+i]*FX[ln[4]];
                FY[ln[i]]-=CD[9*el+5+i]*FY[ln[4]];
                FZ[ln[i]]-=CD[9*el+5+i]*FZ[ln[4]];
            }
        }
        // done_vector(X)
        // done_vector(T)
        done_vector(F)
    }

    /* solving the linear system for each componenet of the vector */
    use_vector(F)
    use_vector(X)

    Vec X1, F1;

    int ierr=VecCreateSeqWithArray(PETSC_COMM_SELF, dim, PETSC_NULL, &X1);
    CHKERR_B(ierr);
    ierr=VecDuplicate(X1, &F1);
    CHKERR_B(ierr);

    /*********************************************************************/
    /*  Solve for the X-coordinate */
    ierr=VecPlaceArray(X1, XX);
    CHKERR_B(ierr);
    ierr=VecPlaceArray(F1, FX);
    CHKERR_B(ierr);

    ret = mat_solve_bm(ksp, sm, bm, F1, X1, dim);

    ierr=VecResetArray(X1);
    CHKERR_B(ierr);
    ierr=VecResetArray(F1);
    CHKERR_B(ierr);

    /*********************************************************************/
    /*  Solve for the Y-coordinate */
    ierr=VecPlaceArray(X1, XY);
    CHKERR_B(ierr);
    ierr=VecPlaceArray(F1, FY);
    CHKERR_B(ierr);

    mat_solve_bm(ksp, sm, bm, F1, X1, dim);

    ierr=VecResetArray(X1);
    CHKERR_B(ierr);
    ierr=VecResetArray(F1);
    CHKERR_B(ierr);

    /*********************************************************************/
    /*  Solve for the Z-coordinate */
    ierr=VecPlaceArray(X1, XZ);
    CHKERR_B(ierr);
    ierr=VecPlaceArray(F1, FZ);
    CHKERR_B(ierr);

    mat_solve_bm(ksp, sm, bm, F1, X1, dim);

    ierr=VecResetArray(X1);
    CHKERR_B(ierr);
    ierr=VecResetArray(F1);
    CHKERR_B(ierr);

    ierr=VecDestroy(X1);
    CHKERR_B(ierr);
    ierr=VecDestroy(F1);
    CHKERR_B(ierr);

    done_vector(X)
    done_vector(F)

        /* post-processing for condensation */
    if(matr->el_type==fe_bub&&matr->condense_data!=NULL&&USE_CONDENSATION)
    {
        int el, *ln;
        double *CD=matr->condense_data;
        use_vector(F)
        use_vector(T)
        use_vector(X)
        for(el=0; el<matr->grid->NEL; el++)
        {
            fe_dofs(matr->grid, el, fe_bub, &ln);
            XX[ln[4]]=FX[ln[4]];
            XY[ln[4]]=FY[ln[4]];
            XZ[ln[4]]=FZ[ln[4]];
            for(i=0; i<4; i++)
            {
                XX[ln[4]]-=CD[9*el+i]*XX[ln[i]];
                XY[ln[4]]-=CD[9*el+i]*XY[ln[i]];
                XZ[ln[4]]-=CD[9*el+i]*XZ[ln[i]];
            }
            XX[ln[4]]/=CD[9*el+4];
            XY[ln[4]]/=CD[9*el+4];
            XZ[ln[4]]/=CD[9*el+4];
        }
        done_vector(X)
        done_vector(T)
        done_vector(F)

    }

    LOG_MPI_WTIME_STOP;
    return ret;
}

void eo_condense_bc_dir(const grid3d_t *grid, unsigned eltype, const int *dofs, unsigned **eos)
{
    eo_bc_dir(grid, eltype, dofs, eos);
    if(eltype==fe_bub)
        *eos[4]=eo_condense;
}

void eo_bc_dir(const grid3d_t *grid, unsigned eltype, const int *dofs, unsigned **eos)
{
    static unsigned _eo[MAX_LDOFS];
    int i, nn;
    if(eos) *eos=_eo;
    else return;
    nn=fe_dofs(grid, 0, eltype, NULL);
    for(i=0; i<nn; i++)
        _eo[i]=fe_dof_bcond(grid, eltype, dofs[i])==BC_DIR?eo_bc:eo_add;
}

void eo_condense_add_all(const grid3d_t *grid, unsigned eltype, const int *dofs, unsigned **eos)
{
    eo_add_all(grid, eltype, dofs, eos);
    if(eltype==fe_bub)
        *eos[4]=eo_condense;
}

void eo_add_all(const grid3d_t *grid, unsigned eltype, const int *dofs, unsigned **eos)
{
    static unsigned _eo[MAX_LDOFS];
    int i, nn;
    (void)dofs;
    if(eos) *eos=_eo;
    else return;
    nn=fe_dofs(grid, 0, eltype, NULL);
    for(i=0; i<nn; i++)
        _eo[i]=eo_add;
}

void fe_matr_print(const fe_matr_p matr, const char *fname, const char *fname_bc)
{
    int ierr;
    PetscViewer file;
    if(fname)
    {
        ierr=PetscViewerASCIIOpen(PETSC_COMM_WORLD, fname, &file);
        CHKERR_B(ierr);
        ierr=MatView(matr->matr, file);
        CHKERR_B(ierr);
        ierr=PetscViewerDestroy(file);
        CHKERR_B(ierr);
    }
    if(fname_bc)
    {
        ierr=PetscViewerASCIIOpen(PETSC_COMM_WORLD, fname_bc, &file);
        CHKERR_B(ierr);
        ierr=MatView(matr->matr_bc, file);
        CHKERR_B(ierr);
        ierr=PetscViewerDestroy(file);
        CHKERR_B(ierr);
    }
}


/****************************************************************************/
#include <memory.h>

void project_grad_vec(const grid3d_t *grid, unsigned type, const fe_matr_p mass,
    vector u, vector gu0, vector gu1, vector gu2,
    vector w, vector t)
{
    assert(mass!=NULL);
    assert(mass->el_type==type);

    int uN=get_N_vec(u);
    scalar c=alloc_scal(uN);

    use_vector(u)
    use_scalar(c)

    // c.F = u->X;
    memcpy(cF, uX, uN*sizeof(double));
    test_grad_e1_in_e2(grid, type, c, type, w);
    fe_matr_solve_vector(mass, gu0, w, t, 0.0);

    // c.F = u->Y;
    memcpy(cF, uY, uN*sizeof(double));
    test_grad_e1_in_e2(grid, type, c, type, w);
    fe_matr_solve_vector(mass, gu1, w, t, 0.0);

    // c.F = u->Z;
    memcpy(cF, uZ, uN*sizeof(double));
    test_grad_e1_in_e2(grid, type, c, type, w);
    fe_matr_solve_vector(mass, gu2, w, t, 0.0);

    done_scalar(c)
    done_vector(u)

    free_scal(c);
    return;
}

void project_grad_scal(const grid3d_t *grid, unsigned type, const fe_matr_p mass,
    scalar p, vector gp, vector w, vector t)
{
    assert(mass!=NULL);
    assert(mass->el_type==type);

    test_grad_e1_in_e2(grid, type, p, type, w);
    fe_matr_solve_vector(mass, w, gp, t, 0.0);

    return;
}


